{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Gotchas of BrainPy Transformations\n",
    "\n",
    "[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/brainpy/brainpy/blob/master/docs_version2/tutorial_FAQs/gotchas_of_brainpy_transforms.ipynb)\n",
    "[![Open in Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/brainpy/brainpy/blob/master/docs_version2/tutorial_FAQs/gotchas_of_brainpy_transforms.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T03:58:51.153959Z",
     "start_time": "2025-10-06T03:58:46.847466Z"
    }
   },
   "source": [
    "import brainpy as bp\n",
    "import brainpy.math as bm\n",
    "\n",
    "bm.set_platform('cpu')\n",
    "\n",
    "bp.__version__"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'3.0.0'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 1
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "BrainPy provides a novel concept for object-oriented transformations based [brainpy.math.Variable](../tutorial_math/variables.ipynb). However, this kind of transformations faces several gotchas:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## 1. ``Variable`` that will be changed cannot be functional arguments"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "This will not work too for the new oo transformations."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T03:58:51.210376Z",
     "start_time": "2025-10-06T03:58:51.158967Z"
    }
   },
   "source": [
    "@bm.jit\n",
    "def f(a, b):\n",
    "  a.value = b\n",
    "\n",
    "a = bm.Variable(bm.ones(1))\n",
    "b = bm.Variable(bm.ones(1) * 10)"
   ],
   "outputs": [],
   "execution_count": 2
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T03:58:53.739860Z",
     "start_time": "2025-10-06T03:58:51.219153Z"
    }
   },
   "source": [
    "f(a, b)\n",
    "\n",
    "try:\n",
    "  assert bm.allclose(a, b)\n",
    "  print('a equals to b.')\n",
    "except:\n",
    "  print('a is not equal to b.')"
   ],
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "Inputs/outputs for brainstate transformations cannot be an instance of State. But we got [1.]",
     "output_type": "error",
     "traceback": [
      "\u001B[31m---------------------------------------------------------------------------\u001B[39m",
      "\u001B[31mJaxStackTraceBeforeTransformation\u001B[39m         Traceback (most recent call last)",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\runpy.py:198\u001B[39m, in \u001B[36m_run_module_as_main\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m    197\u001B[39m     sys.argv[\u001B[32m0\u001B[39m] = mod_spec.origin\n\u001B[32m--> \u001B[39m\u001B[32m198\u001B[39m \u001B[38;5;28;01mreturn\u001B[39;00m _run_code(code, main_globals, \u001B[38;5;28;01mNone\u001B[39;00m,\n\u001B[32m    199\u001B[39m                  \u001B[33m\"\u001B[39m\u001B[33m__main__\u001B[39m\u001B[33m\"\u001B[39m, mod_spec)\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\runpy.py:88\u001B[39m, in \u001B[36m_run_code\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m     81\u001B[39m run_globals.update(\u001B[34m__name__\u001B[39m = mod_name,\n\u001B[32m     82\u001B[39m                    \u001B[34m__file__\u001B[39m = fname,\n\u001B[32m     83\u001B[39m                    __cached__ = cached,\n\u001B[32m   (...)\u001B[39m\u001B[32m     86\u001B[39m                    __package__ = pkg_name,\n\u001B[32m     87\u001B[39m                    __spec__ = mod_spec)\n\u001B[32m---> \u001B[39m\u001B[32m88\u001B[39m exec(code, run_globals)\n\u001B[32m     89\u001B[39m \u001B[38;5;28;01mreturn\u001B[39;00m run_globals\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\ipykernel_launcher.py:18\u001B[39m\n\u001B[32m     16\u001B[39m \u001B[38;5;28;01mfrom\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[34;01mipykernel\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[38;5;28;01mimport\u001B[39;00m kernelapp \u001B[38;5;28;01mas\u001B[39;00m app\n\u001B[32m---> \u001B[39m\u001B[32m18\u001B[39m app.launch_new_instance()\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\traitlets\\config\\application.py:1075\u001B[39m, in \u001B[36mlaunch_instance\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m   1074\u001B[39m app.initialize(argv)\n\u001B[32m-> \u001B[39m\u001B[32m1075\u001B[39m app.start()\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\ipykernel\\kernelapp.py:739\u001B[39m, in \u001B[36mstart\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m    738\u001B[39m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[32m--> \u001B[39m\u001B[32m739\u001B[39m     \u001B[38;5;28mself\u001B[39m.io_loop.start()\n\u001B[32m    740\u001B[39m \u001B[38;5;28;01mexcept\u001B[39;00m \u001B[38;5;167;01mKeyboardInterrupt\u001B[39;00m:\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\tornado\\platform\\asyncio.py:205\u001B[39m, in \u001B[36mstart\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m    204\u001B[39m \u001B[38;5;28;01mdef\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[34mstart\u001B[39m(\u001B[38;5;28mself\u001B[39m) -> \u001B[38;5;28;01mNone\u001B[39;00m:\n\u001B[32m--> \u001B[39m\u001B[32m205\u001B[39m     \u001B[38;5;28mself\u001B[39m.asyncio_loop.run_forever()\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\asyncio\\base_events.py:640\u001B[39m, in \u001B[36mrun_forever\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m    639\u001B[39m \u001B[38;5;28;01mwhile\u001B[39;00m \u001B[38;5;28;01mTrue\u001B[39;00m:\n\u001B[32m--> \u001B[39m\u001B[32m640\u001B[39m     \u001B[38;5;28mself\u001B[39m._run_once()\n\u001B[32m    641\u001B[39m     \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m._stopping:\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\asyncio\\base_events.py:1992\u001B[39m, in \u001B[36m_run_once\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m   1991\u001B[39m     \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[32m-> \u001B[39m\u001B[32m1992\u001B[39m         handle._run()\n\u001B[32m   1993\u001B[39m handle = \u001B[38;5;28;01mNone\u001B[39;00m\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\asyncio\\events.py:88\u001B[39m, in \u001B[36m_run\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m     87\u001B[39m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[32m---> \u001B[39m\u001B[32m88\u001B[39m     \u001B[38;5;28mself\u001B[39m._context.run(\u001B[38;5;28mself\u001B[39m._callback, *\u001B[38;5;28mself\u001B[39m._args)\n\u001B[32m     89\u001B[39m \u001B[38;5;28;01mexcept\u001B[39;00m (\u001B[38;5;167;01mSystemExit\u001B[39;00m, \u001B[38;5;167;01mKeyboardInterrupt\u001B[39;00m):\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\ipykernel\\kernelbase.py:545\u001B[39m, in \u001B[36mdispatch_queue\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m    544\u001B[39m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[32m--> \u001B[39m\u001B[32m545\u001B[39m     \u001B[38;5;28;01mawait\u001B[39;00m \u001B[38;5;28mself\u001B[39m.process_one()\n\u001B[32m    546\u001B[39m \u001B[38;5;28;01mexcept\u001B[39;00m \u001B[38;5;167;01mException\u001B[39;00m:\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\ipykernel\\kernelbase.py:534\u001B[39m, in \u001B[36mprocess_one\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m    533\u001B[39m         \u001B[38;5;28;01mreturn\u001B[39;00m\n\u001B[32m--> \u001B[39m\u001B[32m534\u001B[39m \u001B[38;5;28;01mawait\u001B[39;00m dispatch(*args)\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\ipykernel\\kernelbase.py:437\u001B[39m, in \u001B[36mdispatch_shell\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m    436\u001B[39m     \u001B[38;5;28;01mif\u001B[39;00m inspect.isawaitable(result):\n\u001B[32m--> \u001B[39m\u001B[32m437\u001B[39m         \u001B[38;5;28;01mawait\u001B[39;00m result\n\u001B[32m    438\u001B[39m \u001B[38;5;28;01mexcept\u001B[39;00m \u001B[38;5;167;01mException\u001B[39;00m:\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\ipykernel\\ipkernel.py:362\u001B[39m, in \u001B[36mexecute_request\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m    361\u001B[39m \u001B[38;5;28mself\u001B[39m._associate_new_top_level_threads_with(parent_header)\n\u001B[32m--> \u001B[39m\u001B[32m362\u001B[39m \u001B[38;5;28;01mawait\u001B[39;00m \u001B[38;5;28msuper\u001B[39m().execute_request(stream, ident, parent)\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\ipykernel\\kernelbase.py:778\u001B[39m, in \u001B[36mexecute_request\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m    777\u001B[39m \u001B[38;5;28;01mif\u001B[39;00m inspect.isawaitable(reply_content):\n\u001B[32m--> \u001B[39m\u001B[32m778\u001B[39m     reply_content = \u001B[38;5;28;01mawait\u001B[39;00m reply_content\n\u001B[32m    780\u001B[39m \u001B[38;5;66;03m# Flush output before sending the reply.\u001B[39;00m\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\ipykernel\\ipkernel.py:449\u001B[39m, in \u001B[36mdo_execute\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m    448\u001B[39m \u001B[38;5;28;01mif\u001B[39;00m accepts_params[\u001B[33m\"\u001B[39m\u001B[33mcell_id\u001B[39m\u001B[33m\"\u001B[39m]:\n\u001B[32m--> \u001B[39m\u001B[32m449\u001B[39m     res = shell.run_cell(\n\u001B[32m    450\u001B[39m         code,\n\u001B[32m    451\u001B[39m         store_history=store_history,\n\u001B[32m    452\u001B[39m         silent=silent,\n\u001B[32m    453\u001B[39m         cell_id=cell_id,\n\u001B[32m    454\u001B[39m     )\n\u001B[32m    455\u001B[39m \u001B[38;5;28;01melse\u001B[39;00m:\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\ipykernel\\zmqshell.py:549\u001B[39m, in \u001B[36mrun_cell\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m    548\u001B[39m \u001B[38;5;28mself\u001B[39m._last_traceback = \u001B[38;5;28;01mNone\u001B[39;00m\n\u001B[32m--> \u001B[39m\u001B[32m549\u001B[39m \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28msuper\u001B[39m().run_cell(*args, **kwargs)\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\IPython\\core\\interactiveshell.py:3116\u001B[39m, in \u001B[36mrun_cell\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m   3115\u001B[39m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[32m-> \u001B[39m\u001B[32m3116\u001B[39m     result = \u001B[38;5;28mself\u001B[39m._run_cell(\n\u001B[32m   3117\u001B[39m         raw_cell, store_history, silent, shell_futures, cell_id\n\u001B[32m   3118\u001B[39m     )\n\u001B[32m   3119\u001B[39m \u001B[38;5;28;01mfinally\u001B[39;00m:\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\IPython\\core\\interactiveshell.py:3171\u001B[39m, in \u001B[36m_run_cell\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m   3170\u001B[39m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[32m-> \u001B[39m\u001B[32m3171\u001B[39m     result = runner(coro)\n\u001B[32m   3172\u001B[39m \u001B[38;5;28;01mexcept\u001B[39;00m \u001B[38;5;167;01mBaseException\u001B[39;00m \u001B[38;5;28;01mas\u001B[39;00m e:\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\IPython\\core\\async_helpers.py:128\u001B[39m, in \u001B[36m_pseudo_sync_runner\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m    127\u001B[39m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[32m--> \u001B[39m\u001B[32m128\u001B[39m     coro.send(\u001B[38;5;28;01mNone\u001B[39;00m)\n\u001B[32m    129\u001B[39m \u001B[38;5;28;01mexcept\u001B[39;00m \u001B[38;5;167;01mStopIteration\u001B[39;00m \u001B[38;5;28;01mas\u001B[39;00m exc:\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\IPython\\core\\interactiveshell.py:3394\u001B[39m, in \u001B[36mrun_cell_async\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m   3391\u001B[39m interactivity = \u001B[33m\"\u001B[39m\u001B[33mnone\u001B[39m\u001B[33m\"\u001B[39m \u001B[38;5;28;01mif\u001B[39;00m silent \u001B[38;5;28;01melse\u001B[39;00m \u001B[38;5;28mself\u001B[39m.ast_node_interactivity\n\u001B[32m-> \u001B[39m\u001B[32m3394\u001B[39m has_raised = \u001B[38;5;28;01mawait\u001B[39;00m \u001B[38;5;28mself\u001B[39m.run_ast_nodes(code_ast.body, cell_name,\n\u001B[32m   3395\u001B[39m        interactivity=interactivity, compiler=compiler, result=result)\n\u001B[32m   3397\u001B[39m \u001B[38;5;28mself\u001B[39m.last_execution_succeeded = \u001B[38;5;129;01mnot\u001B[39;00m has_raised\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\IPython\\core\\interactiveshell.py:3639\u001B[39m, in \u001B[36mrun_ast_nodes\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m   3638\u001B[39m     asy = compare(code)\n\u001B[32m-> \u001B[39m\u001B[32m3639\u001B[39m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28;01mawait\u001B[39;00m \u001B[38;5;28mself\u001B[39m.run_code(code, result, async_=asy):\n\u001B[32m   3640\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;01mTrue\u001B[39;00m\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\IPython\\core\\interactiveshell.py:3699\u001B[39m, in \u001B[36mrun_code\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m   3698\u001B[39m     \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[32m-> \u001B[39m\u001B[32m3699\u001B[39m         exec(code_obj, \u001B[38;5;28mself\u001B[39m.user_global_ns, \u001B[38;5;28mself\u001B[39m.user_ns)\n\u001B[32m   3700\u001B[39m \u001B[38;5;28;01mfinally\u001B[39;00m:\n\u001B[32m   3701\u001B[39m     \u001B[38;5;66;03m# Reset our crash handler in place\u001B[39;00m\n",
      "\u001B[36mCell\u001B[39m\u001B[36m \u001B[39m\u001B[32mIn[2]\u001B[39m\u001B[32m, line 5\u001B[39m\n\u001B[32m      3\u001B[39m   a.value = b\n\u001B[32m----> \u001B[39m\u001B[32m5\u001B[39m a = bm.Variable(bm.ones(\u001B[32m1\u001B[39m))\n\u001B[32m      6\u001B[39m b = bm.Variable(bm.ones(\u001B[32m1\u001B[39m) * \u001B[32m10\u001B[39m)\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\BrainPy\\brainpy\\version2\\math\\object_transform\\variables.py:66\u001B[39m, in \u001B[36m__init__\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m     65\u001B[39m Array.\u001B[34m__init__\u001B[39m(\u001B[38;5;28mself\u001B[39m, value, dtype=dtype)\n\u001B[32m---> \u001B[39m\u001B[32m66\u001B[39m brainstate.State.\u001B[34m__init__\u001B[39m(\u001B[38;5;28mself\u001B[39m, value)\n\u001B[32m     68\u001B[39m \u001B[38;5;66;03m# check batch axis\u001B[39;00m\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\_state.py:296\u001B[39m, in \u001B[36m__init__\u001B[39m\u001B[34m()\u001B[39m\n\u001B[32m    292\u001B[39m \u001B[38;5;66;03m# update metadata\u001B[39;00m\n\u001B[32m    293\u001B[39m metadata.update(\n\u001B[32m    294\u001B[39m     _value=value,\n\u001B[32m    295\u001B[39m     _level=_get_trace_stack_level(),\n\u001B[32m--> \u001B[39m\u001B[32m296\u001B[39m     _source_info=source_info_util.current(),\n\u001B[32m    297\u001B[39m     _name=name,\n\u001B[32m    298\u001B[39m     _been_writen=\u001B[38;5;28;01mFalse\u001B[39;00m,\n\u001B[32m    299\u001B[39m     tag=tag,\n\u001B[32m    300\u001B[39m )\n\u001B[32m    302\u001B[39m \u001B[38;5;66;03m# avoid using self._setattr to avoid the check\u001B[39;00m\n",
      "\u001B[31mJaxStackTraceBeforeTransformation\u001B[39m: ValueError: Inputs/outputs for brainstate transformations cannot be an instance of State. But we got [1.]\n\nThe preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.\n\n--------------------",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001B[31mValueError\u001B[39m                                Traceback (most recent call last)",
      "\u001B[36mCell\u001B[39m\u001B[36m \u001B[39m\u001B[32mIn[3]\u001B[39m\u001B[32m, line 1\u001B[39m\n\u001B[32m----> \u001B[39m\u001B[32m1\u001B[39m f(a, b)\n\u001B[32m      3\u001B[39m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[32m      4\u001B[39m   \u001B[38;5;28;01massert\u001B[39;00m bm.allclose(a, b)\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_jit.py:97\u001B[39m, in \u001B[36m_get_jitted_fun.<locals>.jitted_fun\u001B[39m\u001B[34m(*args, **params)\u001B[39m\n\u001B[32m     94\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m fun.fun(*args, **params)\n\u001B[32m     96\u001B[39m \u001B[38;5;66;03m# compile the function and get the state trace\u001B[39;00m\n\u001B[32m---> \u001B[39m\u001B[32m97\u001B[39m state_trace = fun.get_state_trace_by_call(*args, **params)\n\u001B[32m     98\u001B[39m read_state_vals = state_trace.get_read_state_values(\u001B[38;5;28;01mTrue\u001B[39;00m)\n\u001B[32m    100\u001B[39m \u001B[38;5;66;03m# call the jitted function\u001B[39;00m\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_make_jaxpr.py:689\u001B[39m, in \u001B[36mStatefulFunction.get_state_trace_by_call\u001B[39m\u001B[34m(self, compile_if_miss, *args, **kwargs)\u001B[39m\n\u001B[32m    671\u001B[39m \u001B[38;5;28;01mdef\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[34mget_state_trace_by_call\u001B[39m(\u001B[38;5;28mself\u001B[39m, *args, compile_if_miss: \u001B[38;5;28mbool\u001B[39m = \u001B[38;5;28;01mTrue\u001B[39;00m, **kwargs) -> StateTraceStack:\n\u001B[32m    672\u001B[39m \u001B[38;5;250m    \u001B[39m\u001B[33;03m\"\"\"\u001B[39;00m\n\u001B[32m    673\u001B[39m \u001B[33;03m    Read the state trace of the function.\u001B[39;00m\n\u001B[32m    674\u001B[39m \n\u001B[32m   (...)\u001B[39m\u001B[32m    687\u001B[39m \u001B[33;03m        The state trace of the function.\u001B[39;00m\n\u001B[32m    688\u001B[39m \u001B[33;03m    \"\"\"\u001B[39;00m\n\u001B[32m--> \u001B[39m\u001B[32m689\u001B[39m     cache_key = \u001B[38;5;28mself\u001B[39m.get_arg_cache_key(*args, **kwargs, compile_if_miss=compile_if_miss)\n\u001B[32m    690\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m.get_state_trace(cache_key)\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_make_jaxpr.py:839\u001B[39m, in \u001B[36mStatefulFunction.get_arg_cache_key\u001B[39m\u001B[34m(self, compile_if_miss, *args, **kwargs)\u001B[39m\n\u001B[32m    831\u001B[39m cache_key = hashabledict(\n\u001B[32m    832\u001B[39m     static_args=static_args,\n\u001B[32m    833\u001B[39m     dyn_args=dyn_args,\n\u001B[32m    834\u001B[39m     static_kwargs=static_kwargs,\n\u001B[32m    835\u001B[39m     dyn_kwargs=dyn_kwargs,\n\u001B[32m    836\u001B[39m )\n\u001B[32m    838\u001B[39m \u001B[38;5;28;01mif\u001B[39;00m cache_key \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;129;01min\u001B[39;00m \u001B[38;5;28mself\u001B[39m._cached_state_trace \u001B[38;5;129;01mand\u001B[39;00m compile_if_miss:\n\u001B[32m--> \u001B[39m\u001B[32m839\u001B[39m     \u001B[38;5;28mself\u001B[39m.make_jaxpr(*args, **kwargs)\n\u001B[32m    841\u001B[39m \u001B[38;5;28;01mreturn\u001B[39;00m cache_key\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_make_jaxpr.py:913\u001B[39m, in \u001B[36mStatefulFunction.make_jaxpr\u001B[39m\u001B[34m(self, *args, **kwargs)\u001B[39m\n\u001B[32m    888\u001B[39m \u001B[38;5;250m\u001B[39m\u001B[33;03m\"\"\"\u001B[39;00m\n\u001B[32m    889\u001B[39m \u001B[33;03mCreate the JAX Jaxpr representation given example arguments.\u001B[39;00m\n\u001B[32m    890\u001B[39m \n\u001B[32m   (...)\u001B[39m\u001B[32m    909\u001B[39m \u001B[33;03m    If State objects are passed as arguments or returned from the function.\u001B[39;00m\n\u001B[32m    910\u001B[39m \u001B[33;03m\"\"\"\u001B[39;00m\n\u001B[32m    912\u001B[39m \u001B[38;5;66;03m# check input types\u001B[39;00m\n\u001B[32m--> \u001B[39m\u001B[32m913\u001B[39m jax.tree.map(\u001B[38;5;28mself\u001B[39m._check_input_ouput, (args, kwargs), is_leaf=\u001B[38;5;28;01mlambda\u001B[39;00m x: \u001B[38;5;28misinstance\u001B[39m(x, State))\n\u001B[32m    915\u001B[39m \u001B[38;5;66;03m# static args\u001B[39;00m\n\u001B[32m    916\u001B[39m cache_key = \u001B[38;5;28mself\u001B[39m.get_arg_cache_key(*args, **kwargs)\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\jax\\_src\\tree.py:155\u001B[39m, in \u001B[36mmap\u001B[39m\u001B[34m(f, tree, is_leaf, *rest)\u001B[39m\n\u001B[32m    115\u001B[39m \u001B[38;5;28;01mdef\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[34mmap\u001B[39m(f: Callable[..., Any],\n\u001B[32m    116\u001B[39m         tree: Any,\n\u001B[32m    117\u001B[39m         *rest: Any,\n\u001B[32m    118\u001B[39m         is_leaf: Callable[[Any], \u001B[38;5;28mbool\u001B[39m] | \u001B[38;5;28;01mNone\u001B[39;00m = \u001B[38;5;28;01mNone\u001B[39;00m) -> Any:\n\u001B[32m    119\u001B[39m \u001B[38;5;250m  \u001B[39m\u001B[33;03m\"\"\"Maps a multi-input function over pytree args to produce a new pytree.\u001B[39;00m\n\u001B[32m    120\u001B[39m \n\u001B[32m    121\u001B[39m \u001B[33;03m  Args:\u001B[39;00m\n\u001B[32m   (...)\u001B[39m\u001B[32m    153\u001B[39m \u001B[33;03m    - :func:`jax.tree.reduce`\u001B[39;00m\n\u001B[32m    154\u001B[39m \u001B[33;03m  \"\"\"\u001B[39;00m\n\u001B[32m--> \u001B[39m\u001B[32m155\u001B[39m   \u001B[38;5;28;01mreturn\u001B[39;00m tree_util.tree_map(f, tree, *rest, is_leaf=is_leaf)\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\jax\\_src\\tree_util.py:361\u001B[39m, in \u001B[36mtree_map\u001B[39m\u001B[34m(f, tree, is_leaf, *rest)\u001B[39m\n\u001B[32m    359\u001B[39m leaves, treedef = tree_flatten(tree, is_leaf)\n\u001B[32m    360\u001B[39m all_leaves = [leaves] + [treedef.flatten_up_to(r) \u001B[38;5;28;01mfor\u001B[39;00m r \u001B[38;5;129;01min\u001B[39;00m rest]\n\u001B[32m--> \u001B[39m\u001B[32m361\u001B[39m \u001B[38;5;28;01mreturn\u001B[39;00m treedef.unflatten(f(*xs) \u001B[38;5;28;01mfor\u001B[39;00m xs \u001B[38;5;129;01min\u001B[39;00m \u001B[38;5;28mzip\u001B[39m(*all_leaves))\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\jax\\_src\\tree_util.py:361\u001B[39m, in \u001B[36m<genexpr>\u001B[39m\u001B[34m(.0)\u001B[39m\n\u001B[32m    359\u001B[39m leaves, treedef = tree_flatten(tree, is_leaf)\n\u001B[32m    360\u001B[39m all_leaves = [leaves] + [treedef.flatten_up_to(r) \u001B[38;5;28;01mfor\u001B[39;00m r \u001B[38;5;129;01min\u001B[39;00m rest]\n\u001B[32m--> \u001B[39m\u001B[32m361\u001B[39m \u001B[38;5;28;01mreturn\u001B[39;00m treedef.unflatten(f(*xs) \u001B[38;5;28;01mfor\u001B[39;00m xs \u001B[38;5;129;01min\u001B[39;00m \u001B[38;5;28mzip\u001B[39m(*all_leaves))\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_make_jaxpr.py:793\u001B[39m, in \u001B[36mStatefulFunction._check_input_ouput\u001B[39m\u001B[34m(self, x)\u001B[39m\n\u001B[32m    791\u001B[39m \u001B[38;5;28;01mdef\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[34m_check_input_ouput\u001B[39m(\u001B[38;5;28mself\u001B[39m, x):\n\u001B[32m    792\u001B[39m     \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28misinstance\u001B[39m(x, State):\n\u001B[32m--> \u001B[39m\u001B[32m793\u001B[39m         x.raise_error_with_source_info(\n\u001B[32m    794\u001B[39m             \u001B[38;5;167;01mValueError\u001B[39;00m(\n\u001B[32m    795\u001B[39m                 \u001B[33m'\u001B[39m\u001B[33mInputs/outputs for brainstate transformations cannot be an instance of State. \u001B[39m\u001B[33m'\u001B[39m\n\u001B[32m    796\u001B[39m                 \u001B[33mf\u001B[39m\u001B[33m'\u001B[39m\u001B[33mBut we got \u001B[39m\u001B[38;5;132;01m{\u001B[39;00mx\u001B[38;5;132;01m}\u001B[39;00m\u001B[33m'\u001B[39m\n\u001B[32m    797\u001B[39m             )\n\u001B[32m    798\u001B[39m         )\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\_state.py:441\u001B[39m, in \u001B[36mState.raise_error_with_source_info\u001B[39m\u001B[34m(self, error)\u001B[39m\n\u001B[32m    439\u001B[39m name_stack = source_info_util.current_name_stack() + \u001B[38;5;28mself\u001B[39m.source_info.name_stack\n\u001B[32m    440\u001B[39m \u001B[38;5;28;01mwith\u001B[39;00m source_info_util.user_context(\u001B[38;5;28mself\u001B[39m.source_info.traceback, name_stack=name_stack):\n\u001B[32m--> \u001B[39m\u001B[32m441\u001B[39m     \u001B[38;5;28;01mraise\u001B[39;00m error\n",
      "\u001B[31mValueError\u001B[39m: Inputs/outputs for brainstate transformations cannot be an instance of State. But we got [1.]"
     ]
    }
   ],
   "execution_count": 3
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:58:53.742215500Z",
     "start_time": "2023-06-20T12:45:56.284488800Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Variable(value=Array([1.]), dtype=float32)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "All ``Variable``s should be used in a global context.\n",
    "\n",
    "Instead, this works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:58:53.742215500Z",
     "start_time": "2023-06-20T12:45:56.315633300Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "@bm.jit\n",
    "def f(b):\n",
    "  a.value = b\n",
    "\n",
    "a = bm.Variable(bm.ones(1))\n",
    "b = bm.Variable(bm.ones(1) * 10)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:58:53.742215500Z",
     "start_time": "2023-06-20T12:45:56.315633300Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Variable(value=Array([10.]), dtype=float32)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f(b)\n",
    "\n",
    "a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## 2. Functions to be transformed are called twice"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "The core mechanism of any brainpy transformation is that it firsts calls the function to automatically find all ``Variable``s used in the model, and then it calls the function again to compile the model with the found ``Variable``s.\n",
    "\n",
    "Therefore, any function that the user create will be called more than twice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:58:53.742215500Z",
     "start_time": "2023-06-20T12:45:56.331360700Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "@bm.jit\n",
    "def f(inp):\n",
    "  print('calling f ...')\n",
    "  return inp\n",
    "\n",
    "@bm.jit\n",
    "def g(inp):\n",
    "  print('calling g ...')\n",
    "  return f(inp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Taking the above function as an example, when we use this function, we will get:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:58:53.742215500Z",
     "start_time": "2023-06-20T12:45:56.348765200Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "calling g ...\n",
      "calling f ...\n",
      "calling g ...\n",
      "calling f ...\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array(1., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g(1.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "It sequentially calls ``f`` and ``g`` to infer all dynamical variables (instances of ``Variable``) used in these two functions. So we got first two lines of ``calling g ...`` and ``calling f``.\n",
    "\n",
    "Then, it compiles the two functions, so that we got next two lines of ``calling g ...`` and ``calling f``."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Note that this property may get what are not correct in the Python level variables. For example, when we use a global variable to record the number of times the function called:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:58:53.746247900Z",
     "start_time": "2023-06-20T12:45:56.366624400Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "num = [0]\n",
    "\n",
    "@bm.jit\n",
    "def h(inp):\n",
    "  num[0] += 1\n",
    "  return inp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:58:53.746247900Z",
     "start_time": "2023-06-20T12:45:56.381731800Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(1., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h(1.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Although we called the function ``h`` once, we got the number of ``2``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:58:53.746247900Z",
     "start_time": "2023-06-20T12:45:56.397486Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[2]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
